import numpy as np
from scipy.sparse import coo_matrix, csc_matrix, csr_matrix, block_diag
from scipy.sparse import spdiags, eye, bmat, tril, triu
from scipy.sparse.linalg import cg, inv, dsolve,  gmres, minres, LinearOperator, spsolve_triangular
try:
    from mumps import spsolve
except  ImportError:
    from scipy.sparse.linalg import spsolve

import pyamg

from fealpy.functionspace.VectorLagrangeFiniteElementSpace import VectorLagrangeFiniteElementSpace
from fealpy.functionspace.LagrangeFiniteElementSpace import LagrangeFiniteElementSpace
from HuZhangFiniteElementSpace2Ds import HuZhangFiniteElementSpace
from fealpy.fem.integral_alg import IntegralAlg
from fealpy.fem.doperator import stiff_matrix
from boundarycondition import NatureBC
from boundarycondition import EssentialBC
from Inter_Solver import Inter_Solver
from timeit import default_timer as timer
import cProfile

class LinearElasticityFEMModel:
    def __init__(self, mesh,  pde, p, integrator):
        self.mesh = mesh
        self.tensorspace = HuZhangFiniteElementSpace(mesh, p)
        self.vectorspace = VectorLagrangeFiniteElementSpace(mesh, p-1, spacetype='D') 
        self.cspace = LagrangeFiniteElementSpace(mesh, 1) # linear space 
        self.dim = self.tensorspace.dim
        self.sh = self.tensorspace.function()
        self.uh = self.vectorspace.function()
        self.pde = pde
        self.integrator = integrator
        self.measure = mesh.entity_measure()
        self.integralalg = IntegralAlg(self.integrator, mesh, self.measure)
        self.count = 0

        self.itype = mesh.itype
        self.ftype = mesh.ftype



    def get_mass_matrix(self,lam=None):
        'get tensorspace matrix'
        mesh = self.mesh
        tspace = self.tensorspace
        vspace = self.vectorspace
        gdim = tspace.geo_dimension()
        bcs, ws = self.integrator.quadpts, self.integrator.weights

        if gdim == 2:
            d = np.array([1, 1, 2])
        elif gdim == 3:
            d = np.array([1, 1, 1, 2, 2, 2])

        tldof = tspace.number_of_local_dofs()
        tdim = tspace.tensor_dimension()
        NC = self.mesh.number_of_cells()
        NQ = bcs.shape[0]

        phi = tspace.basis(bcs).reshape(NQ,NC,-1,tdim) #(NQ,NC,ldof,tdim)
        aphi = self.pde.compliance_tensor(phi,lam=lam)

        M = np.einsum('i, ijkm, m, ijom, j->jko', ws, aphi, d, phi, self.measure, optimize=True)

        tcell2dof = tspace.cell_to_dof().reshape(NC,-1)
        I = np.einsum('ij, k->ijk', tcell2dof, np.ones(tldof))
        J = I.swapaxes(-1, -2)
        tgdof = tspace.number_of_global_dofs()
        M = csr_matrix((M.flat, (I.flat, J.flat)), shape=(tgdof, tgdof))

        return M



    def get_div_operator_matrix(self):
        mesh = self.mesh
        tspace = self.tensorspace
        vspace = self.vectorspace
        gdim = tspace.geo_dimension()
        bcs, ws = self.integrator.quadpts, self.integrator.weights


        tldof = tspace.number_of_local_dofs()
        tdim = tspace.tensor_dimension()
        NC = self.mesh.number_of_cells()
        NQ = bcs.shape[0]

        
        tgdof = tspace.number_of_global_dofs()
        vgdof = vspace.number_of_global_dofs()
        vldof = vspace.number_of_local_dofs()
        

        dphi = tspace.div_basis(bcs).reshape(NQ,NC,-1,gdim) #(NQ, NC, ldof*tdim, gdim)
        uphi = vspace.basis(bcs) #(NQ,1,vldof,gdim)
        #print(dphi.shape,uphi.shape)
        B = np.einsum('i, ijkm, ijom, j->jko', ws, uphi, dphi, self.measure, optimize=True)

        I = np.einsum('ij, k->ijk', vspace.cell_to_dof(), np.ones(tldof,dtype=int))
        J = np.einsum('ij, k->ikj', tspace.cell_to_dof().reshape(NC,-1,), np.ones(vldof,dtype=int))
        #print(B.shape,I.dtype,J.shape)
        B = csr_matrix((B.flat, (I.flat, J.flat)), shape=(vgdof, tgdof))

        return B


    def get_right_vector(self):
        tspace = self.tensorspace
        vspace = self.vectorspace

        bcs, ws = self.integrator.quadpts, self.integrator.weights
        pp = vspace.mesh.bc_to_point(bcs)
        fval = self.pde.source(pp)#(NQ,NC,gdim)
        phi = vspace.basis(bcs)#(NQ,1,vldof,gdim)
        #print(phi.shape,fval.shape)
        bb = np.einsum('i, ikm, ikjm, k->kj', ws, fval, phi, self.measure)

        cell2dof = vspace.cell_to_dof()
        vgdof = vspace.number_of_global_dofs()
        b = np.bincount(cell2dof.flat, weights=bb.flat, minlength=vgdof)
        return  -b


    def get_diag_mass_matrix(self):
        'get the diagoal of tensorspace matrix for lam=0, mu=0.5'
        tspace = self.tensorspace
        tgdof = tspace.number_of_global_dofs()
        gdim = tspace.geo_dimension()
        bcs, ws = self.integrator.quadpts, self.integrator.weights
        phi = tspace.basis(bcs)

        # construct diag matrix D
        if gdim == 2:
            d = np.array([1, 1, 2])
        elif gdim == 3:
            d = np.array([1, 1, 1, 2, 2, 2])

        D = np.einsum('i, ijklm, m, ijklm, j->jkl', ws, phi, d, phi, self.measure) #(NC,ldof,tdim)
        tcell2dof = tspace.cell_to_dof() #(NC,ldof,tdim)
        D = np.bincount(tcell2dof.flat, weights=D.flat, minlength=tgdof) 

        return D

    def get_mass_matrix_v(self):
        'get the vector sapce matrix'
        vspace = self.vectorspace
        bcs, ws = self.integrator.quadpts, self.integrator.weights
        cellmeasure = self.measure
        vgdof = vspace.number_of_global_dofs()

        vcell2dof=vspace.cell_to_dof()#(NC,vlodf)
        uphi = vspace.basis(bcs) #(NQ,1,vldof,gdim)

        M_u_h = np.einsum('i,ijkl,ijml,j->jkm',ws,uphi,uphi,cellmeasure)#(NC,vldof,vldof)

        I = np.broadcast_to(vcell2dof[:,:,None],shape=M_u_h.shape)#(NC,vldof,vldof)
        J = np.broadcast_to(vcell2dof[:,None,:],shape=M_u_h.shape)#(NC,vldof,vldof)

        M_u_h = csr_matrix((M_u_h.flat,(I.flat,J.flat)),shape=(vgdof,vgdof))
        



        return M_u_h


    def get_A_and_b(self):
        pde = self.pde
        tensorspace = self.tensorspace
        vectorspace = self.vectorspace
        
        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof

        start = timer()
        M = self.get_mass_matrix()
        B = self.get_div_operator_matrix()
        b = self.get_right_vector()
        A = bmat([[M, B.transpose()], [B, None]]).tocsr()
        b = np.r_[np.zeros(tgdof), b]

        #print(np.max(np.abs(A-A.T)),bb), 此时A,b,对称        
        #边界给的是法向或者切向分量，完全体现再gD中
        bc = NatureBC(tensorspace,pde.dirichlet,threshold=pde.is_dirichlet_boundary)
        b = bc.apply(b)
        #print(np.max(np.abs(A-A.T)),bb)   ##此时A,b,对称  
        

        bc = EssentialBC(tensorspace,pde.neumann,threshold=pde.is_neumann_boundary)
        A, b = bc.apply(A,b)

        end = timer()
        print("Construct linear system time:", end - start)

        return A, b


    def Precondieitoner_0(self):
        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        M_h = self.get_mass_matrix(lam=0)
        B_h = self.get_div_operator_matrix()
        S_h = B_h@spsolve(M_h,B_h.transpose())

        M_h_inv = spsolve(M_h,csr_matrix(np.eye(tgdof)))
        S_h_inv = spsolve(S_h,csr_matrix(np.eye(vgdof)))
        P_h = bmat([[M_h_inv,None],[None,S_h_inv]]).tocsr()
        #P_h = spsolve(P_h,csr_matrix(np.eye(tgdof+vgdof)))

        return P_h

    def Precondieitoner_1(self):
        D_h = self.get_diag_mass_matrix()
        #M_h = self.get_mass_matrix(lam=0)
        #print(np.max(np.abs(D_h-2*self.pde.mu*M_h.diagonal())))
        #D_h = M_h.diagonal()
        B_h = self.get_div_operator_matrix()
        S_h = B_h/D_h@B_h.transpose()

        tgdof = len(D_h)
        ii = np.arange(tgdof)
        D_h_inv = csr_matrix((1/D_h,(ii,ii)),shape=(tgdof,tgdof))
        #D_h = csr_matrix((D_h,(ii,ii)),shape=(tgdof,tgdof))
        

        #S_h = B_h@spsolve(D_h,B_h.transpose())
        #print(np.max(np.abs(S_h-B_h@spsolve(D_h,B_h.transpose()))))

        vgdof = S_h.shape[0]
        S_h_inv = spsolve(S_h,np.eye(vgdof))
        

        #M_u_h = self.get_mass_matrix_v()
        #S_h = M_u_h
        #print(type(S_h))

        P_h = bmat([[D_h_inv,None],[None,S_h_inv]]).tocsr()
        #P_h = bmat([[D_h,None],[None,S_h]]).tocsr()

        return P_h

    def Precondieitoner_2(self):
        M_h = self.get_mass_matrix()
        B_h = self.get_div_operator_matrix()
        M_u_h = self.get_mass_matrix_v()

        P_h = bmat([[M_h+B_h.transpose()@spsolve(M_u_h,B_h),None],[None,M_u_h]]).tocsr()

        return P_h

    def Precondieitoner_3(self):
        B_h = self.get_div_operator_matrix()
        D_h = self.get_diag_mass_matrix()
        S_h = B_h/D_h@B_h.transpose()
        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()

        ii = np.arange(tgdof)
        D_h_inv = csr_matrix((1/D_h,(ii,ii)),shape=(tgdof,tgdof))
        P_h_left = bmat([[csr_matrix(np.eye(tgdof)),D_h_inv@B_h.transpose()],
                        [None,-csr_matrix(np.eye(vgdof))]]).tocsr()

        
        D_h = csr_matrix((D_h,(ii,ii)),shape=(tgdof,tgdof))
        P_h_right =spsolve(bmat([[D_h,None],[B_h,S_h]]).tocsr(),csr_matrix(np.eye(tgdof+vgdof)))

        P_h = P_h_left@P_h_right

        return P_h

    def Precondieitoner_4(self):

        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof

        P_h = LinearOperator((gdof, gdof), matvec=self.linear_operator)

        return P_h









    def solve(self):

        #load A and b
        A, b = self.get_A_and_b()

        #load precondieitoner
        start = timer()
        pre_idx = 1
        if pre_idx == 0:
            P_h = self.Precondieitoner_0()
        elif pre_idx == 1:
            P_h = self.Precondieitoner_1()
        elif pre_idx == 2:
            P_h = self.Precondieitoner_2()
        elif pre_idx == 3:
            P_h = self.Precondieitoner_3()

        end = timer()
        print("Construct precondieitioner time:", end-start)

        
        start = timer()
        #x = spsolve(A, bb)
        #tgdof = self.tensorspace.number_of_global_dofs()
        #A = A[:tgdof,:tgdof]
        #P_h = P_h[:tgdof,:tgdof]
        #P_h = self.get_mass_matrix(lam=0)[:tgdof,:tgdof]
        #b = b[:tgdof]
        ISr = Inter_Solver(A,b,maxiter=150,tol=1e-8)
        Solver_idx = 3
        if Solver_idx == 0:
            x = ISr.MINRES()
            r = A@x-b
            print(A.shape)
            print(np.sqrt(r@r)/np.sqrt(b@b))
        elif Solver_idx == 1:
            x = ISr.MINRES_Pre(P_h)
            r = A@x-b
            print(A.shape)
            print(np.sqrt(r@r)/np.sqrt(b@b))
        elif Solver_idx == 2:
            x, exitCode = minres(A,b,maxiter=60,tol=1e-8,M=P_h,show=True)
            #x, exitCode = minres(A,b,maxiter=2000,tol=1e-15,show=True)
            r = A@x-b
            print(np.sqrt(r@r)/np.sqrt(b@b))
        elif Solver_idx == 3:
            
            x, exitCode = gmres(A,b,maxiter=100,tol=1e-8,M=P_h)
            r = A@x-b
            print(np.sqrt(r@r)/np.sqrt(b@b))


        end = timer()
        print("Solve time:", end-start)




        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        self.sh[:] = x[0:tgdof]
        self.uh[:] = x[tgdof:tgdof+vgdof]






    def error(self):
        self.sI = self.tensorspace.interpolation(self.pde.stress)
        self.uI = self.vectorspace.interpolation(self.pde.displacement)

        dim = self.dim
        mesh = self.mesh

        s = self.pde.stress
        sh = self.sh.value 
        e0 = self.integralalg.L2_error(s, sh)

        ds = self.pde.div_stress
        dsh = self.sh.div_value
        e1 = self.integralalg.L2_error(ds, dsh)

        u = self.pde.displacement
        uh = self.uh.value
        e2 = self.integralalg.L2_error(u, uh)

        sI = self.sI.value 
        e3 = self.integralalg.L2_error(s, sI)

        dsI = self.sI.div_value
        e4 = self.integralalg.L2_error(ds, dsI)

        uI = self.uI.value
        e5 = self.integralalg.L2_error(u, uI)

        e6 = self.integralalg.L2_error_uI_uh(uI, uh)

        return e0, e1, e2, e3, e4, e5, e6
